import backtrader as bt
import akshare as ak
import pandas as pd
from datetime import datetime


class TestStrategy(bt.Strategy):
    params = (('p1', 5), ('p2', 10), ('printlog', False))

    def __init__(self):
        self.dataclose = self.datas[0].close
        # 跟踪订单 To keep track of pending orders and buy price/commission
        self.order = None
        self.buyprice = None
        self.buycomm = None

        self.sma1 = bt.indicators.SMA(period=self.p.p1)
        self.sma2 = bt.indicators.SMA(period=self.p.p2)

    def log(self, txt, dt=None, doprint=False):
        if self.params.printlog or doprint:
            dt = dt or self.datas[0].datetime.date(0)
            print(f'{dt.isoformat()}, {txt}')

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            return  # 订单提交和成交当前不做处理

        # 检查订单是否成交
        # 注意，如果现金不够的话，订单会被拒接
        if order.status in [order.Completed]:
            if order.isbuy():
                self.log(f'BUY EXECUTED, Price: {round(order.executed.price, 2)},'
                         f' Cost: {round(order.executed.value, 2)}, Comm: {round(order.executed.comm, 2)}')
                self.buyprice = order.executed.price
                self.buycomm = order.executed.comm
            elif order.issell():
                self.log(f'SELL EXECUTED, Price: {round(order.executed.price, 2)},'
                         f' Cost: {round(order.executed.value, 2)}, Comm: {round(order.executed.comm, 2)}')

            self.bar_executed = len(self)

        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('Order Canceled/Margin/Rejected')

        # 记录没有挂起的订单
        self.order = None

    def notify_trade(self, trade):
        if not trade.isclosed:
            return

        self.log(f'OPERATION PROFIT, GROSS {round(trade.pnl, 2)}, NET {round(trade.pnlcomm, 2)}')

    def next(self):
        # Simply log the closing price of the series from the reference
        self.log(f'Close, {round(self.dataclose[0], 2)}')

        # Check if an order is pending ... if yes, we cannot send a 2nd one
        # 检查是否有挂起的订单，如果有的话，不能再发起一个订单
        if self.order:
            return

        # Check if we are in the market， 检查是否在市场（有持仓）
        if not self.position:
            if self.sma1> self.sma2:
                self.log(f'BUY CREATE, {round(self.dataclose[0], 2)}')
                # Keep track of the created order to avoid a 2nd order
                self.order = self.buy()
        else:
            # 这里注意，Len(self)返回的是当前执行的bar数量，每次next会加1.而Self.bar_executed记录的最后一次交易执行时的bar位置。
            if self.sma1 < self.sma2:
                # SELL, SELL, SELL!!! (with all possible default parameters)
                self.log(f'SELL CREATE, {round(self.dataclose[0], 2)}')

                # Keep track of the created order to avoid a 2nd order
                self.order = self.sell()

    def stop(self):
        self.log('(MA1 Period %2d) (MA2 Period %2d) Ending Value %.2f' %
                 (self.params.p1, self.params.p2, self.broker.getvalue()), doprint=True)


def get_data(code):
    df = ak.stock_zh_a_daily(symbol=code, adjust='qfq')
    df.index = pd.to_datetime(df.date)

    return df


if __name__ == '__main__':
    code = 'sz000001'
    df = get_data(code)
    """
    在进行参数优化时，实例化大脑的时候，有2个与参数优化相关的参数：
    optdatas=True：在处理数据时会采用相对节省时间的方式，进而提高优化速度；
    optreturn=True：在返回回测结果时，为了节省时间，只返回与参数优化最相关的内容（params 和 analyzers），
                    而不会返回参数优化不关心的数据（比如 datas, indicators, observers …等）；
    参数优化是基于 multiprocessing 进行多进程处理数据和分析结果的。
    """
    cerebro = bt.Cerebro(optdatas=True, optreturn=True)
    # cerebro.addstrategy(TestStrategy)
    # Add a strategy
    strats = cerebro.optstrategy(
        TestStrategy,
        p1=range(5, 25),
        p2=range(10, 41))
    start_date = datetime(2020, 1, 1)
    end_date = datetime(2023, 5, 29)
    data = bt.feeds.PandasData(dataname=df, fromdate=start_date, todate=end_date)
    cerebro.adddata(data)
    cerebro.broker.setcash(100000)
    cerebro.addsizer(bt.sizers.FixedSize, stake=100)
    cerebro.broker.setcommission(commission=0.001)

    # 添加分析指标
    # 返回年初至年末的年度收益率
    cerebro.addanalyzer(bt.analyzers.AnnualReturn, _name='_AnnualReturn')
    # 计算最大回撤相关指标
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name='_DrawDown')
    # 计算年化收益
    cerebro.addanalyzer(bt.analyzers.Returns, _name='_Returns', tann=252)
    # 计算年化夏普比率
    cerebro.addanalyzer(bt.analyzers.SharpeRatio_A, _name='_SharpeRatio_A')
    # 返回收益率时序
    cerebro.addanalyzer(bt.analyzers.TimeReturn, _name='_TimeReturn')

    """
    注意：在对于多个标的进行参数优化过程中（比如连续对1000个股票的均线策略寻优），
         如果对于多进程的cpu使用数量不加限制，会有一定几率出现异常错误的情况，
         这类错误目前还没找到解决方法。建议是限制cpu的数量，如设置为2或3：
    """
    result = cerebro.run(maxcpus=3)


    # 打印结果
    def get_my_analyzer(result):
        analyzer = {}
        # 返回参数
        analyzer['period1'] = result.params.p1
        analyzer['period2'] = result.params.p2
        # 提取年化收益
        analyzer['年化收益率'] = result.analyzers._Returns.get_analysis()['rnorm']
        analyzer['年化收益率（%）'] = result.analyzers._Returns.get_analysis()['rnorm100']
        # 提取最大回撤(习惯用负的做大回撤，所以加了负号)
        analyzer['最大回撤（%）'] = result.analyzers._DrawDown.get_analysis()['max']['drawdown'] * (-1)
        # 提取夏普比率
        analyzer['年化夏普比率'] = result.analyzers._SharpeRatio_A.get_analysis()['sharperatio']

        return analyzer


    ret = []
    for i in result:
        ret.append(get_my_analyzer(i[0]))

    pd.DataFrame(ret)
    # cerebro.plot()  # 画图
